#include "ScatterPlot.h"
#include "ScatterVal.h"

#include <variant>
#include <iostream>

#include "TFile.h"
#include "TTree.h"
#include "TSystem.h"
#include "TCanvas.h"
#include "TGraph.h"
#include "TAxis.h"
#include "TLeaf.h"

typedef std::variant<Int_t, Float_t, Double_t> BranchVal;

void DrawScatterPlot(std::string name, YAML::Node config, TChain* chain)
{
    //TGraph
    std::string xName = config["XName"].as<std::string>();
    std::string yName = config["YName"].as<std::string>();

    std::string cut = config["Cut"].as<std::string>("1");
    TFile* tempFile = nullptr;
    TTree* tree = nullptr;
    if(cut == std::string("1"))
    {
        tree = chain;
    }
    else
    {
        tempFile = new TFile("TempFile.root", "recreate");
        tree = chain->CopyTree(cut.c_str());
    }
    Long64_t numOfEntries = tree->GetEntries();

    ScatterVal xArray(xName, tree);
    ScatterVal yArray(yName, tree);

    //Scale
    if(config["Scale"].IsDefined())
    {
        if(config["Scale"]["X"].IsScalar())
        {
            double multi = config["Scale"]["X"].as<double>(1);
            for(Long64_t iX = 0; iX < numOfEntries; ++iX)
            {
                (*xArray.GetVectorPtr())[iX] = multi * (*xArray.GetVectorPtr())[iX];
            }
        }
        else if(config["Scale"]["X"].IsSequence())
        {
            double multi = config["Scale"]["X"][0].as<double>(1);
            double plus  = config["Scale"]["X"][1].as<double>(0);
            for(Long64_t iX = 0; iX < numOfEntries; ++iX)
            {
                (*xArray.GetVectorPtr())[iX] = multi * (*xArray.GetVectorPtr())[iX] + plus;
            }
        }
        if(config["Scale"]["Y"].IsScalar())
        {
            double multi = config["Scale"]["Y"].as<double>(1);
            for(Long64_t iY = 0; iY < numOfEntries; ++iY)
            {
                (*yArray.GetVectorPtr())[iY] = multi * (*yArray.GetVectorPtr())[iY];
            }
        }
        else if(config["Scale"]["Y"].IsSequence())
        {
            double multi = config["Scale"]["Y"][0].as<double>(1);
            double plus  = config["Scale"]["Y"][1].as<double>(0);
            for(Long64_t iY = 0; iY < numOfEntries; ++iY)
            {
                (*yArray.GetVectorPtr())[iY] = multi * (*yArray.GetVectorPtr())[iY] + plus;
            }
        }
    }

    TGraph* graph = new TGraph(numOfEntries, &(*xArray.GetVectorPtr())[0], &(*yArray.GetVectorPtr())[0]);

    //Title
    std::string xTitle = config["XTitle"].as<std::string>(xName);
    std::string yTitle = config["YTitle"].as<std::string>(yName);
    graph->GetXaxis()->SetTitle(xTitle.c_str());
    graph->GetYaxis()->SetTitle(yTitle.c_str());

    //Range
    if(config["XRange"].IsSequence())
        graph->GetXaxis()->SetRangeUser(config["XRange"][0].as<double>(), config["XRange"][1].as<double>());
    if(config["YRange"].IsSequence())
        graph->GetXaxis()->SetRangeUser(config["YRange"][0].as<double>(), config["YRange"][1].as<double>());

    //MarkerStyle
    graph->SetMarkerStyle(config["MarkerStyle"].as<int>(0));

    //Draw
    TCanvas* can = new TCanvas("can", "", 800, 600);
    can->cd();
    graph->Draw(config["DrawOption"].as<std::string>("AP").c_str());

    //Save
    std::stringstream plotStream;
    std::string plotType = config["PlotType"].as<std::string>("eps");
    plotStream << "ezDraw_Scatter_" << name << "." << plotType;
    std::string plotName = config["PlotName"].as<std::string>(plotStream.str());
    can->SaveAs(plotName.c_str());

    
    chain->SetBranchStatus("*", 1);
    if(cut != std::string("1"))
    {
        delete tree;
        delete tempFile;
        gSystem->Exec("rm -f TempFile.root");
    }
    delete graph;
    delete can;
}